import matplotlib.pyplot as plt
import numpy as np
import numpy.linalg as ng
import math

xlst = [i * 0.5 for i in range(21)]
ylst = [2.9, 2.7, 4.8, 5.3, 7.1, 7.6, 7.7, 7.6, 9.4, 9.0, 9.6, 10.0, 10.2, 9.7, 8.3, 8.4, 9.0, 8.3, 6.6, 6.7, 4.1]
N = len(xlst)

# Programming A: Normal equations
# Inner product
def inner(fun1, fun2):
    sum = 0
    for i in range(N):
        sum += fun1(xlst[i]) * fun2(xlst[i])
    return sum

def innery(func):
    sum = 0
    for i in range(N):
        sum += func(xlst[i]) * ylst[i]
    return sum

# Norm
def norm(func):
    return math.sqrt(inner(func, func))

# Initial polynomial
def y1(x):
    return 1

def y2(x):
    return x

def y3(x):
    return (x * x)

# Matrix G
g11 = inner(y1, y1)
g12 = inner(y1, y2)
g13 = inner(y1, y3)
g22 = inner(y2, y2)
g23 = inner(y2, y3)
g33 = inner(y3, y3)
G = np.mat([[g11, g12, g13], [g12, g22, g23], [g13, g23, g33]])
# Matrix c
c = np.mat([[innery(y1)], [innery(y2)], [innery(y3)]])
# Get the result
a = np.dot(ng.inv(G), c)
print("The coefficients by normal equations are :")
print(a)

# Plot the figure
def f(x):
    return float((a[2] * x + a[1]) * x + a[0])

x = np.linspace(-1, 12, 200)
y = []
for t in x:
    yt = f(t)
    y.append(yt)

plt.plot(x, y, label="y = a0 + a1x + a2 x^2")
plt.scatter(xlst, ylst, c="red", s=10)
plt.xlabel("x")
plt.ylabel("y")
plt.title("Least square via normal equations")
plt.legend()
plt.savefig("Figure.png")


# Programming B: QR factorization
x1 = [1 for i in range(N)]
x3 = [s * s for s in xlst]
A = (np.mat([x1, xlst, x3])).T

b = (np.mat([ylst])).T

Q, R = ng.qr(A)
b1 = np.dot(Q.T, b)
c1 = np.dot(Q.T, b)
xa = np.dot(ng.inv(R), c1)
print("The coefficients by QR factorization are :")
print(xa)

# Condition number based on the 2-norm
cond1 = ng.cond(G, 2)
cond2 = ng.cond(R, 2)
print("The condition number of G is :", cond1)
print("The condition number of R is :", cond2)
